{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if TRACES != 'SIMULATED':\n",
    "    scope.adc.bits_per_sample = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import logging\n",
    "import time\n",
    "import numpy as np\n",
    "from tqdm.notebook import tnrange\n",
    "from ecpy.curves import Curve, Point\n",
    "curve = Curve.get_curve('NIST-P256')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_k(bits=256, tries=100):\n",
    "    import random\n",
    "    for i in range(tries):\n",
    "        k = random.getrandbits(bits)\n",
    "        if k < curve.order and k > 0:\n",
    "            return k\n",
    "    raise ValueError(\"Failed to generate a valid random k after %d tries!\" % self.tries)\n",
    "\n",
    "def regularized_k(input_k, bits=256):\n",
    "    \"\"\"Given input k, return the regularized k that the target processes (which the attack will retrieve).\n",
    "    \"\"\"\n",
    "    assert input_k < curve.order\n",
    "    kr = input_k + curve.order\n",
    "    if kr & 2**bits:\n",
    "        kr -= 2**bits\n",
    "    return kr   \n",
    "\n",
    "def input_k(kr, bits=256):\n",
    "    \"\"\"Given the regularized k that the target processes (which the attack will retrieve), return the regularized k that the target will be processing.\n",
    "    \"\"\"\n",
    "    if kr < curve.order:\n",
    "        kr += 2**bits\n",
    "    i_k = kr - curve.order\n",
    "    assert i_k < curve.order # sanity check\n",
    "    return i_k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def new_point():\n",
    "    tries = 100\n",
    "    for i in range(tries):\n",
    "        x = random.getrandbits(256)\n",
    "        y = curve.y_recover(x)\n",
    "        if y:\n",
    "            return (x,y)\n",
    "    raise ValueError('Failed to generate a random point')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ttimes(step=None):\n",
    "    if TRACES == 'SIMULATED':\n",
    "        try:\n",
    "            from cwtraces import sca205_lab_data\n",
    "            load_traces = sca204_lab_data[\"get_traces\"]\n",
    "        except:\n",
    "            print(\"Simulated traces require cwtraces (pip install cwtraces)\")\n",
    "            raise\n",
    "        raws = load_traces('uecc_ttimes%s.npz' % step)\n",
    "        ttimes = raws['arr_0']\n",
    "        raws.close()\n",
    "    else:\n",
    "        ttimes = scope.trigger.get_trigger_times()\n",
    "        if TRACES == 'COLLECT' and step:\n",
    "            np.savez_compressed('data/uecc_ttimes%s.npz' % step, np.asarray(ttimes, dtype=np.uint16))\n",
    "    return ttimes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def capture_ecc_traces(k, N=1, check_sad_triggers=False, check_ttimes=False, Px=None, Py=None, operation=\"pmult\", check=True, get_trace=True, as_int=True, step=1, trim=False):\n",
    "    from chipwhisperer.common.traces import Trace\n",
    "    traces = []\n",
    "    if TRACES == 'SIMULATED':\n",
    "        # eh maybe not optimal but it works\n",
    "        try:\n",
    "            from cwtraces import sca205_lab_data\n",
    "            load_traces = sca204_lab_data[\"get_traces\"]\n",
    "        except:\n",
    "            print(\"Simulated traces require cwtraces (pip install cwtraces)\")\n",
    "            raise\n",
    "        raws = load_traces('uecc_%s.npz' % step)\n",
    "        for t in raws['arr_0']:\n",
    "            traces.append(Trace(t[0], t[1], t[2], None))\n",
    "        raws.close()\n",
    "        print('Pre-recorded traces loaded ✅.')\n",
    "    \n",
    "    else:\n",
    "        if Px:\n",
    "            fixed_point = True\n",
    "        else:\n",
    "            fixed_point = False\n",
    "        for i in tnrange(N, desc='Capturing traces'):\n",
    "            if not fixed_point:\n",
    "                Px, Py = new_point()\n",
    "\n",
    "            start_cycles = scope.adc.trig_count\n",
    "            scope.arm()\n",
    "\n",
    "            if operation == 'pmult':\n",
    "                textout = run_pmult(k, Px, Py, check=check, verbose=False)\n",
    "            else:\n",
    "                logging.error(\"Please supply a valid operation to run.\")\n",
    "\n",
    "            ret = scope.capture()\n",
    "            cycles = scope.adc.trig_count - start_cycles\n",
    "\n",
    "            if ret:\n",
    "                logging.warning(\"Timeout happened during capture\")\n",
    "                next\n",
    "\n",
    "            textin = {'operation': operation,\n",
    "                      'Px': Px,\n",
    "                      'Py': Py,\n",
    "                      'k': k\n",
    "                      }\n",
    "            textout['cycles'] = cycles\n",
    "            if get_trace:\n",
    "                wave = scope.get_last_trace(as_int=as_int)\n",
    "                if trim:\n",
    "                    assert scope.adc.segments > 1, 'trim feature is intended to be used with segmented capture'\n",
    "                    trimmed_wave = np.zeros(scope.adc.segments*trim, dtype=np.uint8)\n",
    "                    #print('XXX %d' % len(trimmed_wave))\n",
    "                    for s in range(scope.adc.segments):\n",
    "                        trimmed_wave[s*trim:(s+1)*trim] = wave[s*scope.adc.samples:s*scope.adc.samples+trim]\n",
    "                    wave = trimmed_wave\n",
    "                    #print('XXX %d' % len(trimmed_wave))\n",
    "                    #print('XXX %d' % len(wave))\n",
    "                skip = False\n",
    "                if check_sad_triggers:\n",
    "                    if scope.SAD.num_triggers_seen not in [255,256]:\n",
    "                        print('bad trace: skipping because scope.SAD.num_triggers_seen = %d' % scope.SAD.num_triggers_seen)\n",
    "                        skip = True\n",
    "                if not skip and check_ttimes:\n",
    "                    ttimes = scope.trigger.get_trigger_times()\n",
    "                    if max(ttimes)/min(ttimes) > 1.15:\n",
    "                        print('bad trace: min(ttimes)=%d, max(ttimes)=%d' % (min(ttimes), max(ttimes)))\n",
    "                        skip = True\n",
    "                    textout['ttimes'] = ttimes\n",
    "                if not skip:\n",
    "                    traces.append(Trace(wave, textin, textout, None))\n",
    "\n",
    "        if TRACES == 'COLLECT':\n",
    "            np.savez_compressed('data/uecc_%s.npz' % step, np.asarray(traces, dtype=object))\n",
    "    \n",
    "    return traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def capture_ecc_trace(k, Px=None, Py=None, operation=\"pmult\", check=True, get_trace=True, as_int=True, step=1):\n",
    "    from chipwhisperer.common.traces import Trace\n",
    "    \n",
    "    if TRACES == 'SIMULATED':\n",
    "        raise ValueError('Not available.')\n",
    "    \n",
    "    else:\n",
    "        start_cycles = scope.adc.trig_count\n",
    "        scope.arm()\n",
    "\n",
    "        if operation == 'pmult':\n",
    "            textout = run_pmult(k, Px, Py, check=check, verbose=False)\n",
    "        else:\n",
    "            logging.error(\"Please supply a valid operation to run.\")\n",
    "\n",
    "        ret = scope.capture()\n",
    "        cycles = scope.adc.trig_count - start_cycles\n",
    "\n",
    "        if ret:\n",
    "            logging.warning(\"Timeout happened during capture\")\n",
    "            return None\n",
    "\n",
    "        textin = {'operation': operation,\n",
    "                  'Px': Px,\n",
    "                  'Py': Py,\n",
    "                  'k': k\n",
    "                  }\n",
    "        textout['cycles'] = cycles\n",
    "        if get_trace:\n",
    "            wave = scope.get_last_trace(as_int=as_int)\n",
    "            trace = Trace(wave, textin, textout, None)\n",
    "        else:\n",
    "            trace = None\n",
    "    \n",
    "    return trace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_pmult(k, Px=None, Py=None, check=True, verbose=False):\n",
    "    \"\"\"Run an arbitrary pmult.\n",
    "    Args:\n",
    "        Px (int): X coordinate of curve point\n",
    "        Py (int): Y coordinate of curve point\n",
    "        k (int): multiplier\n",
    "        check: if set, verify the result (using ecpy)\n",
    "    \"\"\"\n",
    "    from ecpy.curves import Curve,Point\n",
    "    curve = Curve.get_curve('NIST-P256')\n",
    "    if Px == None:\n",
    "        op = 'f'\n",
    "        Px = curve.generator.x\n",
    "        Py = curve.generator.y\n",
    "    else:\n",
    "        op = 'k'\n",
    "        target.simpleserial_write('a', bytearray(int.to_bytes(Px, 32, byteorder='big')))\n",
    "        time.sleep(0.1)\n",
    "        target.simpleserial_write('b', bytearray(int.to_bytes(Py, 32, byteorder='big')))\n",
    "        time.sleep(0.1)\n",
    "\n",
    "    target.simpleserial_write(op, bytearray(int.to_bytes(k, 32, byteorder='big')))\n",
    "\n",
    "    if not target.is_done():\n",
    "        logging.warning (\"Target not done yet, increase clksleeptime!\")\n",
    "        #let's wait a bit more, see what happens:\n",
    "        i = 0\n",
    "        while not target.is_done():\n",
    "            i += 1\n",
    "            time.sleep(0.05)\n",
    "            if i > 100:\n",
    "                logging.warning(\"Target still did not finish operation!\")\n",
    "                break\n",
    "\n",
    "    time.sleep(1)\n",
    "\n",
    "    target.simpleserial_write('p', bytes([0]*32))\n",
    "    Rx = target.simpleserial_read('r', 32)\n",
    "\n",
    "    target.simpleserial_write('q', bytes([0]*32))\n",
    "    Ry = target.simpleserial_read('r', 32)\n",
    "\n",
    "    Rx = int.from_bytes(Rx, byteorder='big')\n",
    "    Ry = int.from_bytes(Ry, byteorder='big')\n",
    "    \n",
    "    # optionally check result:\n",
    "    if check:\n",
    "        P = Point(Px, Py, curve)\n",
    "        Q = k*P\n",
    "        if verbose:\n",
    "            print(\"Expecting Qx = %s\" % hex(Q.x))\n",
    "            print(\"Expecting Qy = %s\" % hex(Q.y))\n",
    "        if Q.x != Rx:\n",
    "            print(\"Bad Rx!\")\n",
    "            print(\"expected %s\" % hex(Q.x))\n",
    "            print(\"got      %s\" % hex(Rx))\n",
    "        if Q.y != Ry:\n",
    "            print(\"Bad Ry!\")\n",
    "            print(\"expected %s\" % hex(Q.y))\n",
    "            print(\"got      %s\" % hex(Ry))\n",
    "    return {'Rx': Rx, 'Ry': Ry}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rerun_pmult(k, verbose=False):\n",
    "    \"\"\"Re-run an arbitrary pmult, using the same point that was used previously.\n",
    "    Args:\n",
    "        k (int): multiplier\n",
    "    \"\"\"\n",
    "    op = 'k'\n",
    "    trace.simpleserial_write(op, bytearray(int.to_bytes(k, 32, byteorder='big')))\n",
    "\n",
    "    if not target.is_done():\n",
    "        logging.warning (\"Target not done yet, increase clksleeptime!\")\n",
    "        #let's wait a bit more, see what happens:\n",
    "        i = 0\n",
    "        while not target.is_done():\n",
    "            i += 1\n",
    "            time.sleep(0.05)\n",
    "            if i > 100:\n",
    "                logging.warning(\"Target still did not finish operation!\")\n",
    "                break\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def capture_ecc_trace_cwlite(k, Px=None, Py=None, operation=\"pmult\", check=True, verbose=False, cycles=6000000):\n",
    "    from chipwhisperer.common.traces import Trace\n",
    "    from tqdm import tnrange\n",
    "    import math\n",
    "    import numpy as np\n",
    "    start_cycles = scope.adc.trig_count\n",
    "\n",
    "    if operation == 'pmult':\n",
    "        from ecpy.curves import Curve,Point\n",
    "        curve = Curve.get_curve('NIST-P256')\n",
    "        if Px == None:\n",
    "            op = 'f'\n",
    "            Px = curve.generator.x\n",
    "            Py = curve.generator.y\n",
    "        else:\n",
    "            op = 'k'\n",
    "            target.simpleserial_write('a', bytearray(int.to_bytes(Px, 32, byteorder='big')))\n",
    "            time.sleep(0.1)\n",
    "            target.simpleserial_write('b', bytearray(int.to_bytes(Py, 32, byteorder='big')))\n",
    "            time.sleep(0.1)\n",
    "\n",
    "        segments = math.ceil(cycles / scope.adc.samples)\n",
    "        scope.adc.offset = 0\n",
    "        wave = np.array([])\n",
    "        for j in tnrange(segments, desc='Capturing trace segments'):\n",
    "            scope.arm()\n",
    "            trace.simpleserial_write(op, bytearray(int.to_bytes(k, 32, byteorder='big')))\n",
    "            #wavesegment = scope.get_last_trace()\n",
    "            while not target.is_done():\n",
    "                pass\n",
    "            time.sleep(1)\n",
    "            ret = scope.capture()\n",
    "            if ret:\n",
    "                print(\"Failed capture\")\n",
    "                continue\n",
    "            wave = np.append(wave, scope.get_last_trace())\n",
    "            scope.adc.offset += scope.adc.samples\n",
    "\n",
    "        time.sleep(1)\n",
    "\n",
    "        target.simpleserial_write('p', bytes([0]*32))\n",
    "        Rx = target.simpleserial_read('r', 32)\n",
    "\n",
    "        target.simpleserial_write('q', bytes([0]*32))\n",
    "        Ry = target.simpleserial_read('r', 32)\n",
    "\n",
    "        Rx = int.from_bytes(Rx, byteorder='big')\n",
    "        Ry = int.from_bytes(Ry, byteorder='big')\n",
    "\n",
    "        # optionally check result:\n",
    "        if check:\n",
    "            P = Point(Px, Py, curve)\n",
    "            Q = k*P\n",
    "            if verbose:\n",
    "                print(\"Expecting Qx = %s\" % hex(Q.x))\n",
    "                print(\"Expecting Qy = %s\" % hex(Q.y))\n",
    "            if Q.x != Rx:\n",
    "                print(\"Bad Rx!\")\n",
    "                print(\"expected %s\" % hex(Q.x))\n",
    "                print(\"got      %s\" % hex(Rx))\n",
    "            if Q.y != Ry:\n",
    "                print(\"Bad Ry!\")\n",
    "                print(\"expected %s\" % hex(Q.y))\n",
    "                print(\"got      %s\" % hex(Ry))\n",
    "        textout = {'Rx': Rx, 'Ry': Ry}        \n",
    "\n",
    "    else:\n",
    "        logging.error(\"Please supply a valid operation to run.\")\n",
    "\n",
    "    cycles = scope.adc.trig_count - start_cycles\n",
    "\n",
    "    textin = {'operation': operation,\n",
    "              'Px': Px,\n",
    "              'Py': Py,\n",
    "              'k': k\n",
    "              }\n",
    "    textout['cycles'] = cycles\n",
    "\n",
    "    if len(wave) >= 1:\n",
    "        return Trace(wave, textin, textout, None)\n",
    "    else:\n",
    "        return None\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
